#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/6/28 9:48
# @Author  : LiShan
# @Email   : lishan_1997@126.com
# @File    : bg_model.py
# @Note    : this is note

import os
from shutil import move
import cv2.cv2 as cv
import time


save_path = "./Test_MOG_FG_BG"
width = 480
height = 360

path_list = [
 save_path,
 save_path + "/fram",
 save_path + "/bg",
 save_path + "/fgmask",
 save_path + "/fg",
 "./backup",
]

for path in path_list:
    if os.path.exists(path):
        pass
    else:
        os.mkdir(path)

def backup_file(source, target):
    # 移动所有模型文件到备份文件夹下
    for root, dirs, files in os.walk(source):
        if len(files) > 0:
            # 获取系统时间
            now = time.strftime("%Y-%m-%d-%H_%M_%S", time.localtime(time.time()))
            # 创建新备份子文件夹
            path = target + now + '/'
            if os.path.exists(path):
                pass
            else:
                os.makedirs(path)
            # 移动文件
            for file in files:
                move(os.path.join(root, file), path)

# TODO 背景减除算法集合
ALGORITHMS_TO_EVALUATE = [
    (cv.bgsegm.createBackgroundSubtractorGMG(20, 0.7), 'GMG'),
    (cv.bgsegm.createBackgroundSubtractorCNT(), 'CNT'),
    (cv.createBackgroundSubtractorKNN(detectShadows=False), 'KNN'),
    (cv.bgsegm.createBackgroundSubtractorMOG(), 'MOG'),
    (cv.createBackgroundSubtractorMOG2(detectShadows=False), 'MOG2'),
    (cv.bgsegm.createBackgroundSubtractorGSOC(), 'GSOC'),
    (cv.bgsegm.createBackgroundSubtractorLSBP(), 'LSBP'),
]


# TODO 主函数
def main():
    backup_file(save_path + "/", "./backup/")
    # 背景分割识别器序号
    # algo_index = 2   # KNN
    algo_index = 4    # MOG2
    subtractor = ALGORITHMS_TO_EVALUATE[algo_index][0]
    videoPath = "./video/background_model.mp4"
    show_fgmask = False

    # 获得运行环境CPU的核心数
    nthreads = cv.getNumberOfCPUs()
    # 设置线程数
    cv.setNumThreads(nthreads)

    # 读取视频
    capture = cv.VideoCapture(videoPath)

    # 当前帧数
    frame_num = 0
    # 总执行时间
    sum_Time = 0.0

    while True:
        ret, frame = capture.read()
        if not ret:
            return
        frame = cv.resize(frame, (width, height))
        frame_num = frame_num + 1
        # begin_time = time()
        # gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
        model = frame.copy()
        fgmask = subtractor.apply(model)
        bg = subtractor.getBackgroundImage(model)
        # end_time = time()
        # run_time = end_time - begin_time
        # sum_Time = sum_Time + run_time
        # # 平均执行时间
        # average_Time = sum_Time / (frame_num + 1)
        #
        fg = frame.copy()
        if show_fgmask:
            segm = fgmask
        else:
            segm = (fg * 0.5).astype('uint8')
            cv.add(fg, (100, 100, 0, 0), segm, fgmask)

        if frame_num % 25 == 0 or frame_num == 1:
            # fg = cv.bitwise_and(fg, fg, mask=fgmask)
            cv.imwrite(save_path + "/fram/" + str(frame_num) + "_frame" + ".jpg", frame)
            cv.imwrite(save_path + "/fgmask/" + str(frame_num) + "_fgmask" + ".jpg", fgmask)
            cv.imwrite(save_path + "/bg/" + str(frame_num) + "_bg" + ".jpg", bg)
            cv.imwrite(save_path + "/fg/" + str(frame_num) + "_fg" + ".jpg", segm)
        #
        # # 显示当前方法
        # cv.putText(segm, ALGORITHMS_TO_EVALUATE[algo_index][1], (10, 30), cv.FONT_HERSHEY_PLAIN, 2.0, (255, 0, 255),
        #             2,
        #             cv.LINE_AA)
        # # 显示当前线程数
        # cv.putText(segm, str(nthreads) + " threads", (10, 60), cv.FONT_HERSHEY_PLAIN, 2.0, (255, 0, 255), 2,
        #             cv.LINE_AA)
        # # 显示当前每帧执行时间
        # cv.putText(segm, "averageTime {} s".format(average_Time), (10, 90), cv.FONT_HERSHEY_PLAIN, 2.0,
        #             (255, 0, 255), 2, cv.LINE_AA)
        #
        # cv.imshow('some', segm)
        key = cv.waitKey(1) & 0xFF
        # 按Esc退出循环
        if key == 27:
            break

    cv.destroyAllWindows()


if __name__ == '__main__':
    main()

